#对于彩色图像分离各频段使用了另一种表示图像的方式
import numpy as np
import cv2
import matplotlib.pyplot as plt
def normalize_to_uint8(img):
    img = img - img.min()
    img = img / (img.max() + 1e-8)  # 防止除零
    return (img * 255).astype(np.uint8)


def frequency_masks(shape, low_ratio=0.1, high_ratio=0.3):
    H, W = shape
    center_u, center_v = H // 2, W // 2
    U, V = np.ogrid[:H, :W]
    dist = np.sqrt((U - center_u)**2 + (V - center_v)**2)
    max_dist = dist.max()
    low_mask = dist <= low_ratio * max_dist
    mid_mask = (dist > low_ratio * max_dist) & (dist <= high_ratio * max_dist)
    high_mask = dist > high_ratio * max_dist
    return low_mask, mid_mask, high_mask

def process_y_channel_bands(img_path):
    img_bgr = cv2.imread(img_path)
    img_ycc = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCrCb)
    y, cr, cb = cv2.split(img_ycc)
    y = y.astype(np.float32)

    H, W = y.shape
    low_mask, mid_mask, high_mask = frequency_masks((H, W))

    # DFT
    F = np.fft.fft2(y)
    F_shifted = np.fft.fftshift(F)

    F_low = F_shifted * low_mask
    F_mid = F_shifted * mid_mask
    F_high = F_shifted * high_mask

    y_low = np.fft.ifft2(np.fft.ifftshift(F_low)).real
    y_mid = np.fft.ifft2(np.fft.ifftshift(F_mid)).real
    y_high = np.fft.ifft2(np.fft.ifftshift(F_high)).real


    def merge_and_show(y_new, cr, cb, title, enhance=False):
        if enhance:
            y_new = normalize_to_uint8(y_new)
        else:
            y_new = np.clip(y_new, 0, 255).astype(np.uint8)

        img_ycc_new = cv2.merge([y_new, cr, cb])
        img_rgb = cv2.cvtColor(img_ycc_new, cv2.COLOR_YCrCb2RGB)
        plt.imshow(img_rgb);
        plt.title(title);
        plt.axis('off')

    F_combined = F_shifted * (low_mask + high_mask)
    y_combined = np.fft.ifft2(np.fft.ifftshift(F_combined)).real
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 黑体字体，Windows系统常见
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    plt.figure(figsize=(12, 6))
    plt.subplot(2, 3, 1)
    merge_and_show(y, cr, cb, "原图")
    plt.subplot(2, 3, 2)
    merge_and_show(y_low, cr, cb, "低频")
    plt.subplot(2, 3, 3)
    merge_and_show(y_high, cr, cb, "高频")
    plt.subplot(2, 3, 5)
    merge_and_show(y_combined, cr, cb, "保留低频+高频", enhance=True)
    plt.subplot(2, 3, 6)
    merge_and_show(y_mid, cr, cb, "中频（对比）", enhance=True)
    plt.tight_layout()
    plt.show()

# 替换路径
process_y_channel_bands('D:\\2D-DFT\\image\\test.jpg')
